// "Copyright [2021] <Copyright Shengkai Lin>"

#ifndef PROTOCOLS_ROCE_ROCE_ADAPTER_H_
#define PROTOCOLS_ROCE_ROCE_ADAPTER_H_

#include <functional>
#include <unordered_map>
#include <utility>

#include "inet/applications/udpapp/UdpBasicApp.h"
#include "inet/transportlayer/contract/udp/UdpSocket.h"
#include "ai_sim.h"
#include "RoceConnectionBase.h"

using namespace ai_sim;

namespace roce {

class RoceAdapter : public inet::UdpBasicApp {
  public:
    RoceAdapter() {}

  protected:
    void handleMessageWhenUp(omnetpp::cMessage *msg) override;
    void processStart() override;

  private:
    using SrcDstPair = std::pair<int, int>;
    struct SrcDstPairHash {
        std::size_t operator() (const SrcDstPair& pair) const {
            return std::hash<int>()(pair.first) ^ std::hash<int>()(pair.second);
        }
    };

    void CheckAndSend(int dst);
    void HandleUpperLayerPacket(FlowMsg* f_msg);
    void HandleLowerLayerPacket(inet::Packet* packet);

    inet::L3Address GetL3Address(int idx);
    std::unique_ptr<RoceConnectionBase>& FindOrCreateRoceConnection(int dst_idx);

    std::unordered_map<int, std::unique_ptr<RoceConnectionBase>>
        dst_to_roce_connections_;
    
    RoceType roce_type_;
    int local_idx_;
    int local_port_;
    int dst_port_;
    bool dont_fragment_;
    inet::UdpSocket socket_;
};

}  // namespace ai_sim

#endif  // PROTOCOLS_ROCE_ROCE_ADAPTER_H_
